{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9a6dd449",
   "metadata": {},
   "source": [
    "# Lesson 4: Tool Use and Conversational Chess"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75bcfd3c",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6032c971-8934-4018-8743-eace0553578c",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": [
    "llm_config = {\"model\": \"gpt-4-turbo\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "086cb574-1299-43ce-90b4-79e089a76a11",
   "metadata": {
    "height": 63
   },
   "outputs": [],
   "source": [
    "import chess\n",
    "import chess.svg\n",
    "from typing_extensions import Annotated"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4406540",
   "metadata": {},
   "source": [
    "## Initialize the chess board"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10161c38-0fe6-4a03-a5ff-069e8374a1d8",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "board = chess.Board()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c799e0a9-be62-4c20-b524-7a39d125f46d",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "made_move = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22d38427",
   "metadata": {},
   "source": [
    "## Define the needed tools\n",
    "\n",
    "### 1. Tool for getting legal moves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78607e9f-47b0-419e-b652-dd6b5b0390c1",
   "metadata": {
    "height": 114
   },
   "outputs": [],
   "source": [
    "def get_legal_moves(\n",
    "    \n",
    ") -> Annotated[str, \"A list of legal moves in UCI format\"]:\n",
    "    return \"Possible moves are: \" + \",\".join(\n",
    "        [str(move) for move in board.legal_moves]\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0be25bc8",
   "metadata": {},
   "source": [
    "### 2. Tool for making a move on the board"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08a09854-dfd4-495c-84c1-8ea1dd8cacdb",
   "metadata": {
    "height": 505
   },
   "outputs": [],
   "source": [
    "def make_move(\n",
    "    move: Annotated[str, \"A move in UCI format.\"]\n",
    ") -> Annotated[str, \"Result of the move.\"]:\n",
    "    move = chess.Move.from_uci(move)\n",
    "    board.push_uci(str(move))\n",
    "    global made_move\n",
    "    made_move = True\n",
    "    \n",
    "    # Display the board.\n",
    "    display(\n",
    "        chess.svg.board(\n",
    "            board,\n",
    "            arrows=[(move.from_square, move.to_square)],\n",
    "            fill={move.from_square: \"gray\"},\n",
    "            size=200\n",
    "        )\n",
    "    )\n",
    "    \n",
    "    # Get the piece name.\n",
    "    piece = board.piece_at(move.to_square)\n",
    "    piece_symbol = piece.unicode_symbol()\n",
    "    piece_name = (\n",
    "        chess.piece_name(piece.piece_type).capitalize()\n",
    "        if piece_symbol.isupper()\n",
    "        else chess.piece_name(piece.piece_type)\n",
    "    )\n",
    "    return f\"Moved {piece_name} ({piece_symbol}) from \"\\\n",
    "    f\"{chess.SQUARE_NAMES[move.from_square]} to \"\\\n",
    "    f\"{chess.SQUARE_NAMES[move.to_square]}.\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a06e6f41",
   "metadata": {},
   "source": [
    "## Create agents\n",
    "\n",
    "You will create the player agents and a board proxy agents for the chess board."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3055679f-5f69-4d32-aff6-076925a36174",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "from autogen import ConversableAgent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9466949c-c2cf-4bba-8307-d7ff4b70a948",
   "metadata": {
    "height": 165
   },
   "outputs": [],
   "source": [
    "# Player white agent\n",
    "player_white = ConversableAgent(\n",
    "    name=\"Player White\",\n",
    "    system_message=\"You are a chess player and you play as white. \"\n",
    "    \"First call get_legal_moves(), to get a list of legal moves. \"\n",
    "    \"Then call make_move(move) to make a move.\",\n",
    "    llm_config=llm_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcb6a474-2691-425c-8558-beb3363cd2bb",
   "metadata": {
    "height": 165
   },
   "outputs": [],
   "source": [
    "# Player black agent\n",
    "player_black = ConversableAgent(\n",
    "    name=\"Player Black\",\n",
    "    system_message=\"You are a chess player and you play as black. \"\n",
    "    \"First call get_legal_moves(), to get a list of legal moves. \"\n",
    "    \"Then call make_move(move) to make a move.\",\n",
    "    llm_config=llm_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b69ac38-6399-425a-9681-4a2b7403a4bf",
   "metadata": {
    "height": 148
   },
   "outputs": [],
   "source": [
    "def check_made_move(msg):\n",
    "    global made_move\n",
    "    if made_move:\n",
    "        made_move = False\n",
    "        return True\n",
    "    else:\n",
    "        return False\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3eaa260-422c-4e44-a355-42f0d885e9bc",
   "metadata": {
    "height": 131
   },
   "outputs": [],
   "source": [
    "board_proxy = ConversableAgent(\n",
    "    name=\"Board Proxy\",\n",
    "    llm_config=False,\n",
    "    is_termination_msg=check_made_move,\n",
    "    default_auto_reply=\"Please make a move.\",\n",
    "    human_input_mode=\"NEVER\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cd1c45d",
   "metadata": {},
   "source": [
    "## Register the tools\n",
    "\n",
    "A tool must be registered for the agent that calls the tool and the agent that executes the tool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5caf8ae1-4b39-4be5-89fa-0d57b4595841",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "from autogen import register_function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "880a9926-d76b-4322-872c-b691d933506e",
   "metadata": {
    "height": 284
   },
   "outputs": [],
   "source": [
    "for caller in [player_white, player_black]:\n",
    "    register_function(\n",
    "        get_legal_moves,\n",
    "        caller=caller,\n",
    "        executor=board_proxy,\n",
    "        name=\"get_legal_moves\",\n",
    "        description=\"Get legal moves.\",\n",
    "    )\n",
    "    \n",
    "    register_function(\n",
    "        make_move,\n",
    "        caller=caller,\n",
    "        executor=board_proxy,\n",
    "        name=\"make_move\",\n",
    "        description=\"Call this tool to make a move.\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd487d3e-9226-4764-bd25-b452c1993388",
   "metadata": {
    "height": 29
   },
   "outputs": [],
   "source": [
    "player_black.llm_config[\"tools\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a7fa1f0",
   "metadata": {},
   "source": [
    "## Register the nested chats\n",
    "\n",
    "Each player agent will have a nested chat with the board proxy agent to\n",
    "make moves on the chess board."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3230f945-f302-4f39-bc37-766ca082fb28",
   "metadata": {
    "height": 369
   },
   "outputs": [],
   "source": [
    "player_white.register_nested_chats(\n",
    "    trigger=player_black,\n",
    "    chat_queue=[\n",
    "        {\n",
    "            \"sender\": board_proxy,\n",
    "            \"recipient\": player_white,\n",
    "            \"summary_method\": \"last_msg\",\n",
    "        }\n",
    "    ],\n",
    ")\n",
    "\n",
    "player_black.register_nested_chats(\n",
    "    trigger=player_white,\n",
    "    chat_queue=[\n",
    "        {\n",
    "            \"sender\": board_proxy,\n",
    "            \"recipient\": player_black,\n",
    "            \"summary_method\": \"last_msg\",\n",
    "        }\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d371ce02",
   "metadata": {},
   "source": [
    "## Start the Game\n",
    "\n",
    "The game will start with the first message."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ceb2b874",
   "metadata": {},
   "source": [
    "<p style=\"background-color:#ECECEC; padding:15px; \"> <b>Note:</b> In this lesson, you will use GPT 4 for better results. Please note that the lesson has a quota limit. If you want to explore the code in this lesson further, we recommend trying it locally with your own API key."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f445154",
   "metadata": {},
   "source": [
    "**Note**: You might get a slightly different moves than what's shown in the video."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5595758f-7ae1-467f-bad9-79c6c7cd0b25",
   "metadata": {
    "height": 131
   },
   "outputs": [],
   "source": [
    "board = chess.Board()\n",
    "\n",
    "chat_result = player_black.initiate_chat(\n",
    "    player_white,\n",
    "    message=\"Let's play chess! Your move.\",\n",
    "    max_turns=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad9cf403",
   "metadata": {},
   "source": [
    "## Adding a fun chitchat to the game!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc6c2ef8-4807-4b9e-9eff-0da219ac6381",
   "metadata": {
    "height": 165
   },
   "outputs": [],
   "source": [
    "player_white = ConversableAgent(\n",
    "    name=\"Player White\",\n",
    "    system_message=\"You are a chess player and you play as white. \"\n",
    "    \"First call get_legal_moves(), to get a list of legal moves. \"\n",
    "    \"Then call make_move(move) to make a move. \"\n",
    "    \"After a move is made, chitchat to make the game fun.\",\n",
    "    llm_config=llm_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1da2db96-83d1-4b75-aa05-3840913a46ac",
   "metadata": {
    "height": 165
   },
   "outputs": [],
   "source": [
    "player_black = ConversableAgent(\n",
    "    name=\"Player Black\",\n",
    "    system_message=\"You are a chess player and you play as black. \"\n",
    "    \"First call get_legal_moves(), to get a list of legal moves. \"\n",
    "    \"Then call make_move(move) to make a move. \"\n",
    "    \"After a move is made, chitchat to make the game fun.\",\n",
    "    llm_config=llm_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "489e5258-a48d-4745-9cf8-6d530ce1162c",
   "metadata": {
    "height": 692
   },
   "outputs": [],
   "source": [
    "for caller in [player_white, player_black]:\n",
    "    register_function(\n",
    "        get_legal_moves,\n",
    "        caller=caller,\n",
    "        executor=board_proxy,\n",
    "        name=\"get_legal_moves\",\n",
    "        description=\"Get legal moves.\",\n",
    "    )\n",
    "\n",
    "    register_function(\n",
    "        make_move,\n",
    "        caller=caller,\n",
    "        executor=board_proxy,\n",
    "        name=\"make_move\",\n",
    "        description=\"Call this tool to make a move.\",\n",
    "    )\n",
    "\n",
    "player_white.register_nested_chats(\n",
    "    trigger=player_black,\n",
    "    chat_queue=[\n",
    "        {\n",
    "            \"sender\": board_proxy,\n",
    "            \"recipient\": player_white,\n",
    "            \"summary_method\": \"last_msg\",\n",
    "            \"silent\": True,\n",
    "        }\n",
    "    ],\n",
    ")\n",
    "\n",
    "player_black.register_nested_chats(\n",
    "    trigger=player_white,\n",
    "    chat_queue=[\n",
    "        {\n",
    "            \"sender\": board_proxy,\n",
    "            \"recipient\": player_black,\n",
    "            \"summary_method\": \"last_msg\",\n",
    "            \"silent\": True,\n",
    "        }\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "723edd10-e0c6-4205-8e9b-db4cef7a464f",
   "metadata": {
    "height": 131
   },
   "outputs": [],
   "source": [
    "board = chess.Board()\n",
    "\n",
    "chat_result = player_black.initiate_chat(\n",
    "    player_white,\n",
    "    message=\"Let's play chess! Your move.\",\n",
    "    max_turns=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e8ca3c3",
   "metadata": {},
   "source": [
    "**Note:** \n",
    "To add human input to this game, add **human_input_mode=\"ALWAYS\"** for both player agents."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
